//
// Created by mio on 25-4-23.
//

#include "RtmpConnection.h"

#include <RtmpServer.h>
#include <RtmpSession.h>

#include "rtmp.h"

RtmpConnection::RtmpConnection(std::shared_ptr<RtmpServer> rtmp_server, TaskScheduler *scheduler, int sockfd)
    :RtmpConnection(scheduler,sockfd,(Rtmp*)rtmp_server.get())
{
    handshake_.reset(new RtmpHandshake(RtmpHandshake::HANDSHAKE_C0C1));
    rtmp_server_ = rtmp_server;
}

RtmpConnection::RtmpConnection(TaskScheduler* scheduler, int sockfd, Rtmp* rtmp)
    : TcpConnection(scheduler, sockfd),
    rtmp_chunk_(new RtmpChunk()),
    state_(HANDSHAKE)
{
    peer_width_ = rtmp->GetPeerBandwidth();
    acknowledgement_size_ = rtmp->GetAcknowledgementSize();
    max_chunk_size_ = rtmp->GetChunkSize();
    stream_path_ = rtmp->GetStreamPath();
    stream_name_ = rtmp->GetStreamName();
    app_ = rtmp->GetApp();

    this->SetReadCallback([this](std::shared_ptr<TcpConnection> conn, BufferReader& buffer)
    {
        return this->OnRead(buffer);
    });

    this->SetCloseCallback([this](std::shared_ptr<TcpConnection> conn)
    {
        this->OnClose();
    });
}

RtmpConnection::~RtmpConnection()
{
}

bool RtmpConnection::OnRead(BufferReader& buffer)
{
    bool ret = true;
    if (handshake_->IsCompleted()) // 是否握手完成, 完成之后才能发送消息
    {
        ret = HandleChunk(buffer);
    }
    else
    {
        std::shared_ptr<char> res(new char[4096], std::default_delete<char[]>());
        int res_size = handshake_->Parse(buffer, res.get(), 4096);
        if (res_size < 0) ret = false;
        if (res_size > 0) this->Send(res.get(), res_size);
        if (handshake_->IsCompleted())
        {
            // 如果握手刚完成，并且缓冲区中还有剩余数据，尝试处理这些数据
            if (buffer.ReadableBytes() > 0)
                ret = HandleChunk(buffer);
        }
    }
    return ret;
}

void RtmpConnection::OnClose()
{
    this->HandleDeleteStream();
}

bool RtmpConnection::HandleChunk(BufferReader& buffer)
{
    int ret = -1;
    do
    {
        RtmpMessage rtmp_msg;
        ret = rtmp_chunk_->Parse(buffer, rtmp_msg);
        if (ret >= 0) // 解析成功, 需要处理message
        {
            if (rtmp_msg.IsCompleted())
            {
                if (!HandleMessage(rtmp_msg)) return false;
            }
            if (ret == 0) break; // 缓冲区没有数据
        }
        else return false;
    } while (buffer.ReadableBytes() > 0);
    return true;
}

bool RtmpConnection::HandleMessage(RtmpMessage& rtmp_msg)
{
    bool ret = true;
    switch (rtmp_msg.type_id)
    {
    case RTMP_VIDEO:    // 视频消息 (Type ID 9)
        ret = HandleVideo(rtmp_msg);
        break;
    case RTMP_AUDIO:    // 音频消息 (Type ID 8)
        ret = HandleAudio(rtmp_msg);
        break;
    case RTMP_INVOKE:   // 命令消息 (Type ID 20 for AMF0)
        ret = HandleInvoke(rtmp_msg);
        break;
    case RTMP_NOTIFY:   // 数据消息 (Type ID 18 for AMF0)
        ret = HandleNotify(rtmp_msg);
        break;
    case RTMP_SET_CHUNK_SIZE:
        rtmp_chunk_->SetInChunkSize(ReadUint32BE(rtmp_msg.payload.get()));
        break;
    default:
        break;
    }
    return ret;
}

bool RtmpConnection::HandleInvoke(RtmpMessage& rtmp_msg)
{
    bool ret = true;
    amf_decoder_.reset();

    int bytes_used = amf_decoder_.decode(rtmp_msg.payload.get(), rtmp_msg.length, 1);
    if (bytes_used < 0) return false;

    std::string method = amf_decoder_.getString();
    // 需要判断流是否创建
    if (rtmp_msg.stream_id == 0)
    {
        // 需要继续解码后续参数（事务 ID, 命令对象等）
        // 这里的 decode 调用没有指定解码项数，会尝试解码所有剩余数据
        bytes_used += amf_decoder_.decode(rtmp_msg.payload.get() + bytes_used, rtmp_msg.length - bytes_used);
        if (method == "connect") ret = HandleConnect();
        else if (method == "createStream") ret = HandleCreateStream();
    }
    else if (rtmp_msg.stream_id == stream_id_)
    {
        // 解码事务 ID (通常为 0), 命令对象 (null), 和流名称
        // 这里解码 3 项：事务ID(Number), 命令对象(Null), 流名称(String)
        bytes_used += amf_decoder_.decode(rtmp_msg.payload.get() + bytes_used, rtmp_msg.length - bytes_used, 3);
        // 获取流名称并构建流路径
        stream_name_ = amf_decoder_.getString();
        stream_path_ = "/" + app_ + "/" +stream_name_;

        // 如果消息中还有更多参数（例如 play 命令的可选参数），继续解码
        if (rtmp_msg.length > bytes_used)
        {
            bytes_used += amf_decoder_.decode(rtmp_msg.payload.get() + bytes_used, rtmp_msg.length - bytes_used);
        }

        if (method == "publish") ret = HandlePublish();
        else if (method == "play") ret = HandlePlay();
        else if (method == "DeleteStream") ret = HandleDeleteStream();
    }

    return ret;
}

bool RtmpConnection::HandleNotify(RtmpMessage& rtmp_msg)
{
    // 准备解码器
    amf_decoder_.reset();

    // 解码第一个 AMF 项，应该是数据消息的名称 (String)
    int bytes_used = amf_decoder_.decode(rtmp_msg.payload.get(), rtmp_msg.length, 1);
    std::string method = amf_decoder_.getString();
    if (method == "@setDataFrame") // 检查是否是特殊的 @setDataFrame，通常用于标记 onMetaData
    {
        amf_decoder_.reset();
        // 解码 @setDataFrame 后面的第一个参数，应该是 "onMetaData" 字符串
        bytes_used += amf_decoder_.decode(rtmp_msg.payload.get() + bytes_used, rtmp_msg.length - bytes_used, 1);
        if (bytes_used < 0) return false;

        if (amf_decoder_.getString() == "onMetaData")
        {
            // 解码 onMetaData 后面跟着的实际元数据对象 (通常是 ECMA Array)
            // 这里直接调用 decode 解码剩余所有数据
            amf_decoder_.decode(rtmp_msg.payload.get() + bytes_used, rtmp_msg.length - bytes_used);
            meta_data_ = amf_decoder_.getObjects();
        }

        // 设置元数据
        auto server = rtmp_server_.lock();
        if (!server) return false;
        auto session = rtmp_session_.lock();
        if (session) session->SendMetaData(meta_data_);
    }
    return true;
}

bool RtmpConnection::HandleAudio(RtmpMessage& rtmp_msg)
{
    uint8_t type = RTMP_AUDIO;
    uint8_t* payload = (uint8_t*)rtmp_msg.payload.get();
    uint32_t length = rtmp_msg.length;

    // 从第一个字节解析音频格式和编码 ID
    uint8_t sound_format = (payload[0] >> 4) & 0x0f; // 高 4 位: 音频格式 (10=AAC)
    uint8_t code_id = payload[0] & 0x0f; // 编码id

    auto server = rtmp_server_.lock();
    if (!server) return false;
    auto session = rtmp_session_.lock();
    if (!session) return false;

    // 检查是否是 AAC 序列头
    if (sound_format == RTMP_CODEC_ID_AAC && payload[1] == 0) // 说明编码器为aac, 第二字节为0说明这个音频数据为AAC序列头
    {
        aac_sequence_header_size_ = rtmp_msg.length;
        aac_sequence_header_.reset(new char[rtmp_msg.length], std::default_delete<char[]>());
        memcpy(aac_sequence_header_.get(), rtmp_msg.payload.get(), aac_sequence_header_size_);

        // 通知 Session 更新序列头
        session->SetAacSequenceHeader(aac_sequence_header_, aac_sequence_header_size_);
        type = RTMP_AAC_SEQUENCE_HEADER;
    }

    session->SendMediaData(type, rtmp_msg.timestamp, rtmp_msg.payload, rtmp_msg.length);
    return true;
}

bool RtmpConnection::HandleVideo(RtmpMessage& rtmp_msg)
{
    uint8_t type = RTMP_VIDEO;
    uint8_t* payload = (uint8_t*)rtmp_msg.payload.get();
    uint32_t length = rtmp_msg.length;

    // 从第一个字节解析帧类型和编码 ID
    uint8_t frame_type = (payload[0] >> 4) & 0x0f; // 获取帧类型 高 4 位: 1=Keyframe, 2=Interframe, 3=Disposable Interframe, 4=Generated Keyframe, 5=Video info/command frame
    uint8_t code_id = payload[0] & 0x0f; // 视频编码ID 低 4 位: 7=AVC(H.264)

    auto server = rtmp_server_.lock();
    if (!server) return false;
    auto session = rtmp_session_.lock();
    if (!session) return false;

    if (frame_type == 1 && code_id == RTMP_CODEC_ID_H264 && payload[1] == 0) // 更新264序列头
    {
        // 缓存 AVC 序列头
        avc_sequence_header_size_ = rtmp_msg.length;
        avc_sequence_header_.reset(new char[rtmp_msg.length], std::default_delete<char[]>());
        memcpy(avc_sequence_header_.get(), rtmp_msg.payload.get(), avc_sequence_header_size_);

        session->SetAvcSequenceHeader(avc_sequence_header_, avc_sequence_header_size_);
        type = RTMP_AVC_SEQUENCE_HEADER;
    }
    // 将视频数据（无论是普通帧还是序列头）发送给 Session 处理
    session->SendMediaData(type, rtmp_msg.timestamp, rtmp_msg.payload, rtmp_msg.length);
    return true;
}

bool RtmpConnection::HandleConnect()
{
    // 检查 Command Object 中是否包含 "app" 属性
    if (!amf_decoder_.hasObject("app"))
    {
        return false; // 缺少必需的 app 参数，连接失败
    }

    AmfObject amfobj = amf_decoder_.getObject("app");
    app_ = amfobj.amf_string; // 获取引用程序的值
    if (app_ == "") return false;

    /* --- 发送标准的服务器初始化消息 --- */
    SendAcknowledgement();
    SetPeerBandWidth();
    SetChunkSize();

    /* --- 构造并发送 connect 的 _result 响应 --- */
    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("_result", 7);
    amf_encoder_.encodeNumber(amf_decoder_.getNumber()); // 获取 createStream 请求的 Transaction ID
    // 编码 Properties 对象
    objects["fmsVer"] = AmfObject(std::string("FMS/4,5,0,297"));
    objects["capabilities"] = AmfObject(255.0);
    objects["mode"] = AmfObject(1.0);
    amf_encoder_.encodeObjects(objects);
    //清空对象
    objects.clear();
    // 编码 Information 对象
    objects["level"] = AmfObject(std::string("status"));
    objects["code"] = AmfObject(std::string("NetConnection.Connect.Success"));
    objects["description"] = AmfObject(std::string("Connection succeeded"));
    objects["objectEncoding"] = AmfObject(0.0);
    amf_encoder_.encodeObjects(objects);
    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID,amf_encoder_.data(),amf_encoder_.size());
    printf("HandleConnect\n");
    return true;
}

bool RtmpConnection::HandleCreateStream()
{
    int stream_id = rtmp_chunk_->GetStreamId();
    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("_result", 7);
    amf_encoder_.encodeNumber(amf_decoder_.getNumber()); // 获取 createStream 请求的 Transaction ID

    amf_encoder_.encodeObjects(objects);
    amf_encoder_.encodeNumber(stream_id);
    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data(), amf_encoder_.size());
    stream_id_ = stream_id;
    return true;
}

bool RtmpConnection::HandlePublish()
{
    auto server = rtmp_server_.lock();
    if (!server) return false;

    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);                   // 事务 ID 0
    amf_encoder_.encodeObjects(objects);
    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data(), amf_encoder_.size());

    bool is_error = false;

    if (server->HasPublisher(stream_path_)) // 检查服务器是否已经有发布者发布到同一个流路径
    {
        is_error = true;
        objects["level"] = AmfObject(std::string("error"));
        objects["code"] = AmfObject(std::string("NetStream.Publish.BadName")); // 说明当该流已经推送
        objects["description"] = AmfObject(std::string("Stream already publishing."));
    }
    else if (state_ == START_PUBLISH) // 检查当前连接是否已经处于发布状态 (防止重复 publish)
    {
        is_error = true;
        objects["level"] = AmfObject(std::string("error"));
        objects["code"] = AmfObject(std::string("NetStream.Publish.BadConnection")); // 说明当该流已经推送
        objects["description"] = AmfObject(std::string("Stream already publishing."));
    }
    else // 如果没有错误，准备发送成功状态
    {
        objects["level"] = AmfObject(std::string("status"));
        objects["code"] = AmfObject(std::string("NetStream.Publish.Start")); // 说明当该流已经推送
        objects["description"] = AmfObject(std::string("Start publishing."));

        server->AddSession(stream_path_);
        rtmp_session_ = server->GetSession(stream_path_);
        if (server)
        {
            server->NotifyEvent("publish.start", stream_path_);
        }
    }
    amf_encoder_.encodeObjects(objects);
    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data(), amf_encoder_.size());

    if (is_error) return false;
    else
    {
        state_ = START_PUBLISH;
        is_publishing_ = true;
    }
    auto session = rtmp_session_.lock();
    if (session)
    {
        session->AddSink(std::dynamic_pointer_cast<RtmpSink>(shared_from_this()));
    }
    return true;
}

bool RtmpConnection::HandlePlay()
{
    auto server = rtmp_server_.lock();
    if (!server) return false;
    AmfObjects objects;
    amf_encoder_.reset();

    // 发送 onStatus NetStream.Play.Reset
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);                   // 事务 ID 0
    objects["level"] = AmfObject(std::string("status"));
    objects["code"] = AmfObject(std::string("NetStream.Play.Reset"));
    objects["description"] = AmfObject(std::string("Resetting ond playing stream"));
    amf_encoder_.encodeObjects(objects);

    if (!SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data(), amf_encoder_.size()))
        return false;

    // 再发送 onStatus NetStream.Play.Start
    objects.clear();
    amf_encoder_.reset();
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);
    objects["level"] = AmfObject(std::string("status"));
    objects["code"] = AmfObject(std::string("NetStream.Play.Start"));
    objects["description"] = AmfObject(std::string("Started playing."));
    amf_encoder_.encodeObjects(objects);
    if (!SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data(), amf_encoder_.size()))
    {
        return false;
    }
    // 再通知客户端权限 发送 |RtmpSampleAccess 数据消息
    // 这是 Flash Player 播放 H.264/AAC 等需要直接访问样本数据时必需的权限通知
    amf_encoder_.reset();
    amf_encoder_.encodeString("|RtmpSampleAccess", 17);
    amf_encoder_.encodeBoolean(true); // 允许读
    amf_encoder_.encodeBoolean(true); // 允许写
    if (!this->SendNotifyMessage(RTMP_CHUNK_DATA_ID, amf_encoder_.data(), amf_encoder_.size()))
    {
        return false;
    }

    // 更新状态
    state_ = START_PLAY;
    rtmp_session_ = server->GetSession(stream_path_);
    auto session = rtmp_session_.lock();
    if (session) session->AddSink(std::dynamic_pointer_cast<RtmpSink>(shared_from_this()));
    if(server) server->NotifyEvent("Play.start",stream_path_);

    return true;
}

bool RtmpConnection::HandleDeleteStream()
{
    auto server = rtmp_server_.lock();
    if (!server) return false;
    // 检查是否有有效的流路径（意味着之前成功 publish 或 play 过）
    if (stream_path_ != "")
    {
        auto session = rtmp_session_.lock();
        if (!session) return false;
        auto conn = std::dynamic_pointer_cast<RtmpConnection>(shared_from_this());
        GetTaskScheduler()->AddTimer([session, conn]()
        {
            session->RemoveSink(conn);
            return false;
        }, 1);
        if (is_publishing_)
        {
            server->NotifyEvent("publish,stop", stream_path_);
        }
        else if (is_playing_)
        {
            server->NotifyEvent("play.stop", stream_path_);
        }
        is_playing_ = false;
        is_publishing_ = false;
        has_key_frame = false;
        rtmp_chunk_->Clear();
    }

    return true;
}

void RtmpConnection::SetPeerBandWidth()
{
    std::shared_ptr<char> data(new char[5], std::default_delete<char[]>());
    WriteUint32BE(data.get(), peer_width_);
    // 填充第5个字节
    // 0 代表客户端必须按照带宽进行处理, 否则会丢弃
    // 1 是软限制, 可以按照带宽进行处理, 允许超出
    // 2 是动态限制, 表示客户端可以根据网络去调整带宽
    data.get()[4] = 2;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_BANDWIDTH_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 5;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

void RtmpConnection::SendAcknowledgement()
{
    std::shared_ptr<char> data(new char[4], std::default_delete<char[]>());
    WriteUint32BE(data.get(), acknowledgement_size_);

    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_ACK_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 4;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

void RtmpConnection::SetChunkSize()
{
    rtmp_chunk_->SetOutChunkSize(max_chunk_size_);
    std::shared_ptr<char> data(new char[4], std::default_delete<char[]>());
    WriteUint32BE(data.get(), max_chunk_size_);
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_SET_CHUNK_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 4;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

bool RtmpConnection::SendInvokeMessage(uint32_t csid, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (this->IsClosed()) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_INVOKE;
    rtmp_msg.timestamp = 0;
    rtmp_msg.length = payload_size;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    SendRtmpChunks(csid, rtmp_msg);
    return true;
}

bool RtmpConnection::SendNotifyMessage(uint32_t csid, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (this->IsClosed()) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_NOTIFY;
    rtmp_msg.timestamp = 0;
    rtmp_msg.length = payload_size;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    SendRtmpChunks(csid, rtmp_msg);
    return true;
}

bool RtmpConnection::IsKeyFrame(std::shared_ptr<char> data, uint32_t size)
{
    // 从第一个字节获取帧类型和编码 ID
    uint8_t frame_type = (data.get()[0] >> 4) & 0x0f;
    uint8_t code_id = data.get()[0] & 0x0f;
    return (frame_type == 1 && code_id == RTMP_CODEC_ID_H264);
}

void RtmpConnection::SendRtmpChunks(uint32_t csid, RtmpMessage& rtmp_msg)
{
    // 粗略估算需要的缓冲区大小：消息长度 + 每个块可能的头开销 + 额外余量
    uint32_t capacity = rtmp_msg.length + rtmp_msg.length / max_chunk_size_ * 5 + 1024;
    std::shared_ptr<char> buffer(new char[capacity], std::default_delete<char[]>());
    int size = rtmp_chunk_->CreateChunk(csid, rtmp_msg, buffer.get(), capacity);
    if (size > 0)
    {
        this->Send(buffer.get(), size);
    }
}

bool RtmpConnection::SendMetaData(AmfObjects metaData)
{
    if (this->IsClosed()) return false;
    if (metaData.size() == 0) return false;
    amf_encoder_.reset();
    amf_encoder_.encodeString("onMetaData", 10);
    amf_encoder_.encodeECMA(meta_data_);
    if (!SendNotifyMessage(RTMP_CHUNK_DATA_ID, amf_encoder_.data(), amf_encoder_.size()))
    {
        return false;
    }
    return true;
}

bool RtmpConnection::SendMediaData(uint8_t type, uint64_t timestamp, std::shared_ptr<char> payload,
    uint32_t payload_size)
{
    if (this->IsClosed()) return false;
    if (payload_size == 0) return false;
    is_playing_ = true;
    if (type == RTMP_AVC_SEQUENCE_HEADER)
    {
        avc_sequence_header_ = payload;
        avc_sequence_header_size_ = payload_size;
    }
    else if (type == RTMP_AAC_SEQUENCE_HEADER)
    {
        aac_sequence_header_ = payload;
        aac_sequence_header_size_ = payload_size;
    }
    if (!has_key_frame && avc_sequence_header_size_ > 0
        && (type != RTMP_AVC_SEQUENCE_HEADER)
        && (type != RTMP_AAC_SEQUENCE_HEADER)) // 数据包既不是序列头，还没有收到关键帧
    {
        // 开始判断是否为关键帧
        if (IsKeyFrame(payload, payload_size))
        {
            has_key_frame = true;
            // TODO: 发送GOP
        }
        else return true;
    }

    // 收到关键帧之后开始发送message
    RtmpMessage rtmp_msg;
    rtmp_msg.timestamp = timestamp;

    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.length = payload_size;
    rtmp_msg.payload = payload;

    if (type == RTMP_VIDEO || type == RTMP_AVC_SEQUENCE_HEADER)
    {
        rtmp_msg.type_id = RTMP_VIDEO;
        SendRtmpChunks(RTMP_CHUNK_VIDEO_ID, rtmp_msg);
    }
    else if (type == RTMP_AUDIO || type == RTMP_AAC_SEQUENCE_HEADER)
    {
        rtmp_msg.type_id = RTMP_AUDIO;
        SendRtmpChunks(RTMP_CHUNK_AUDIO_ID, rtmp_msg);
    }
    return true;
}
